package chapter1;

/**
 * union-find算法
 */
public abstract class UF {

    private int[] ids;
    private int count;

    public UF(int count) {
        this.count = count;
        ids = new int[this.count];
        for (int i = 0; i < this.count; i++) {
            ids[i] = i;
        }
    }

    public int count() {
        return count;
    }

    public abstract int find(int p);


    public abstract void union(int p, int q);


    public boolean connected(int p, int q) {
        return find(p) == find(q);
    }

    public class QuickFindUF extends UF{

        public QuickFindUF(int count) {
            super(count);
        }

        @Override
        public int find(int p) {
            return ids[p];
        }

        @Override
        public void union(int p, int q) {
            int pID = find(p);
            int qID = find(q);

            if (pID == qID) return;

            for (int i = 0; i < ids.length; i++) {
                if (ids[i] == pID){
                    ids[i] = qID;
                }
            }
            count--;
        }
    }

    public class QuickUnionUF extends UF{

        public QuickUnionUF(int count) {
            super(count);
        }

        @Override
        public int find(int p) {
            while (ids[p] != p){
                p = ids[p];
            }
            return p;
        }

        @Override
        public void union(int p, int q) {
            int pRoot = find(p);
            int qRoot = find(q);
            if (pRoot == qRoot) return;

            ids[p] = qRoot;
            count--;
        }
    }

    /**
     * 带权重的quick-union
     */
    class WeightedQuickUnionUF extends UF{

        private int[] sz;//存储根节点的节点数目

        public WeightedQuickUnionUF(int count) {
            super(count);
            sz = new int[count];
            for (int i = 0; i < count; i++) {
                sz[i] = 1; //初始都一样，都是1
            }
        }

        @Override
        public int find(int p) {
            while (p != ids[p]){
                p = ids[p];
            }
            return p;
        }

        @Override
        public void union(int p, int q) {
            int pRoot = find(p);
            int qRoot = find(q);
            if (pRoot == qRoot){
                return;
            }
            //将小树的根节点，连接到大树的根节点
            if (sz[pRoot] < sz[qRoot]){
                ids[pRoot] = qRoot;
                // 大树根节点数目 += 小树根节点数目
                sz[qRoot] += sz[pRoot];
            }else {
                ids[qRoot] = pRoot;
                sz[pRoot] += sz[qRoot];
            }
            count--;

        }
    }
}
